from dataclasses import dataclass
from typing import Dict, List, Tuple
from enum import Enum
from ..trains.locomotive_dynamics import *


# @njit  # (cache=True)
def point_of_reach(s1, v1, s2, v2):
    """
    交汇点坐标为x； s1, v1; s2, v2分别为两个机车前车和后车的位置和速度

    x = s1 + v1 * t = s2 + v2 * t
    """
    if v2-v1 <= 0:
        return -100.0
    return s1+v1*(s1-s2)/(v2-v1)


# @jitclass
@dataclass
class Block:
    """
    闭塞区间
    """
    start: float
    end: float
    has_station: bool

    @property
    def position(self):
        return (self.start+self.end)/2

    # def __init__(self, start, end):
    #     self.start = start
    #     self.end = end


ALLOW_TRAIN_DOWN_ONLY = 1  # 只允许下行
ALLOW_TRAIN_UP_ONLY = 2  # 只允许上行
ALLOW_TRAIN_BIDIRECTIONAL = 3  # 允许同时上行和下行


# @jitclass(
#     spec=[
#         ("direction", types.int64),
#         ("blocks", types.ListType(Block.class_type.instance_type)),
#         ("train_occupations", types.DictType(
#             types.int64, Train.class_type.instance_type)),
#     ]
# )
class Railroad:

    def __init__(self, block_num: int = 100) -> None:
        self.direction = ALLOW_TRAIN_UP_ONLY
        self.blocks = [Block(i*1000, (i+1)*1000, i % 12 == 0)
                       for i in range(block_num)]
        # 从区间号，映射到区间中的火车
        self.train_occupations: Dict[int, Train] = {}
        # 获取待避的火车数量
        self.train_waiting: Dict[int, List[Train]] = {}

    def find_speed_mismatch(self):
        """
        寻找速度不匹配的两列火车，并且决定是否要安排会让

        """
        # 获取在正线上的火车
        trains: List[Tuple[int, Train]] = [
            (k, v) for k, v in self.train_occupations.items()]

        # 获取所有待避中的火车
        for block_id, trains_waiting in self.train_waiting.items():
            trains.extend([(block_id, train) for train in trains_waiting])

        # 将所有正线上和待避中的火车进行排序
        trains.sort(lambda item: item[0])  # 上行线路上，按照block_id的升序排列各个区间
        print(self.train_occupations, trains)
        for i in range(len(trains)-1):
            train_prev = trains[i][1]
            train_next = trains[i+1][1]
            if not train_prev.waiting:  # 如果前车已经在待避，显然不用计算这个了
                por = point_of_reach(train_prev._distance/1000.0, train_prev.max_speed,
                                     train_next._distance/1000.0, train_next.max_speed)
                distance_to_por = por - train_prev._distance/1000.0  # 到超车点的距离
                print(train_next.max_speed, train_next.waiting, self.free_blocks_up(
                    train_next.current_block_id), self.free_blocks_down(train_next.current_block_id))
                if train_next.waiting and (distance_to_por >= 25.0 or distance_to_por < 0) and self.free_blocks_up(train_next.current_block_id) >= 1 and train_next.current_block_id not in self.train_occupations:
                    train_next.cancel_wait(self)

                if 0.0 < distance_to_por < 25.0:  # 如果快要超车了，前车就要进站待避
                    for block_id in range(train_prev.current_block_id, -1, -1):
                        block = self.blocks[block_id]
                        if block.has_station:
                            train_prev.request_wait(block)
                            break
            # if train_next.waiting

    def put_train(self, train: Train, direction: int) -> int:
        """将火车添加到这段轨道, 并返回在block中的索引"""
        if direction == DIRECTION_DOWN:
            assert self.direction in (
                ALLOW_TRAIN_BIDIRECTIONAL, ALLOW_TRAIN_DOWN_ONLY)
            self.train_occupations[0] = train
            train.current_block_id = 0
            return 0
        elif direction == DIRECTION_UP:
            assert self.direction in (
                ALLOW_TRAIN_BIDIRECTIONAL, ALLOW_TRAIN_UP_ONLY)
            last_index = len(self.blocks)-1
            self.train_occupations[last_index] = train
            train.current_block_id = last_index
            return last_index
        else:
            raise ValueError(direction)

    def set_train_waiting(self, position: int):
        """将这段区间的火车设置为等待状态"""
        if position not in self.train_waiting:
            self.train_waiting[position] = []
        self.train_waiting[position].append(
            self.train_occupations.pop(position))

    def set_train_running(self, train: Train, position: int):
        """将这段区间的某一个火车设置为运行状态"""
        index = self.train_waiting[position].index(train)
        if position in self.train_occupations.keys():
            raise ValueError("Current block was occupied!")
        self.train_occupations[position] = self.train_waiting[position].pop(
            index)
        print("running!", self.train_occupations)

    def remove_train_from_block(self, position: int):
        """将火车从这段轨道移除"""
        self.train_occupations.pop(position)

    def move_train_block(self, original: int, new: int):
        """
        将某个区间中的火车移动到另一个区间
        这个方法由火车类调用
        """
        self.train_occupations[new] = self.train_occupations.pop(original)

    @property
    def length(self):
        return self.blocks[-1].end

    def free_blocks_down(self, current_block_id: int) -> int:
        """
        下行方向的区间状态

        n = 0: 前方区间有车
        n > 0: 前方的n个区间没有车。n最大为3，代表至少3个区间无车
        """
        if current_block_id <= 0:  # 如果已经是最后一个区间，就要返回0
            return 0
        count = 0
        for block_id in range(current_block_id-1, -1, -1):  # 查询后面有几个区间空闲
            # 如果区间无车占用，则继续找，查看下个区间有没有问题
            if block_id not in self.train_occupations:
                count += 1
                if count >= 3:
                    break
            # 如果区间有车占用，直接退出循环
            else:
                break
        return count

    def free_blocks_up(self, current_block_id: int) -> int:
        """
        上行方向的区间状态

        n = 0: 前方区间有车
        n > 0: 前方的n个区间没有车。n最大为3，代表至少3个区间无车
        """
        if current_block_id >= len(self.blocks)-1:  # 如果已经是最后一个区间，就要返回0
            return 0
        count = 0
        for block_id in range(current_block_id+1, len(self.blocks)):  # 查询后面有几个区间空闲
            # 如果区间无车占用，则继续找，查看下个区间有没有问题
            if block_id not in self.train_occupations:
                count += 1
                if count >= 3:
                    break
            # 如果区间有车占用，直接退出循环
            else:
                break
        return count

    def add_train(self, train: "Train"):
        pass
        # self.train_occupations[0] = train

    def update(self, step: int):
        if step % 10 == 0:
            self.find_speed_mismatch()

# class Train(LocomotiveDynamics):


class MyModel:
    def __init__(self) -> None:
        self.train = Train()
        self.rail = Railroad()
        self.rail.add_train(self.train)

    def run(self):
        pass
